﻿using System;
using System.IO;
using Microsoft.Cci;
using Microsoft.Cci.MutableCodeModel;
using Microsoft.Cci.Contracts;
using CSharpSourceEmitter;
using Microsoft.Cci.ILToCodeModel;


namespace Mutator
{
    class Program
    {
        public static SourceEmitterOutputString sourceEmitterOutput;

        static void Main(string[] args)
        {
            if (args == null || args.Length == 0)
            {
                Console.WriteLine("usage: PeToPe [path]fileName.ext [decompile]");
                return;
            }

            HostEnvironment host = new HostEnvironment();
            IModule/*?*/ module = host.LoadUnitFrom(args[0]) as IModule;
            if (module == null || module == Dummy.Module || module == Dummy.Assembly)
            {
                Console.WriteLine(args[0] + " is not a PE file containing a CLR module or assembly.");
                return;
            }
            generate(module, host);
        }
        public static void generate(IModule module, HostEnvironment host)
        {

            Console.WriteLine("      ****************************************************************");
            Console.WriteLine("      *                    PexMutator Started!                       *");
            Console.WriteLine("      *                                                              *");
            Console.WriteLine("      ****************************************************************");
            PdbReader/*?*/ pdbReader = null;
            PdbWriter/*?*/ pdbWriter = null;


            string pdbFile = Path.ChangeExtension(module.Location, "pdb");
            if (File.Exists(pdbFile))
            {
                Stream pdbStream = File.OpenRead(pdbFile);
                pdbReader = new PdbReader(pdbStream, host);
            }

            using (pdbReader)
            {

                //Construct a Code Model from the Metadata model via decompilation
                module = Decompiler.GetCodeModelFromMetadataModel(host, module, pdbReader);

                Mutants mutator;

                ExpressionMutantGenerator expressionMutantGenerator;
                ConditionMutantGenerator conditionMutantGenerator;
                ExpressionMutantRetriever expressionMutantRetriever;

                OriginalVisitor originalVisitor;


                mutator = new Mutants(
                  host, pdbReader);

                expressionMutantGenerator = new ExpressionMutantGenerator(
                  host, pdbReader);
                expressionMutantRetriever = new ExpressionMutantRetriever(
                 host, pdbReader);
                conditionMutantGenerator = new ConditionMutantGenerator(
                  host, pdbReader);
                originalVisitor = new OriginalVisitor(
                  host, pdbReader);


                mutator.expressionMutantGenerator = expressionMutantGenerator;
                mutator.conditionMutantGenerator = conditionMutantGenerator;
                mutator.expressionMutantRetriever = expressionMutantRetriever;
                mutator.originalVisitor = originalVisitor;
                //mutator.sourceVisitor = sourceVisitor;


                module = mutator.Visit(module);


                Stream peStream = File.Create(module.Location.Substring(0,module.Location.LastIndexOf("\\"))+"\\Mutated_" + module.ModuleName);
                if (pdbReader == null)
                {
                    PeWriter.WritePeToStream(module, host, peStream);
                }
                else
                {
                    using (pdbWriter = new PdbWriter(module.Location + ".pdb", pdbReader))
                    {
                        PeWriter.WritePeToStream(module, host, peStream, pdbReader, pdbReader, pdbWriter);
                    }
                }
                Console.WriteLine();
                Console.WriteLine("****** Successfully Finished!");
                Console.WriteLine("****** Generated Mutant Number: " + Mutants.mutantNumber);
                Console.WriteLine("****** Generated Meta-Program: " + module.Location.Substring(0, module.Location.LastIndexOf("\\")) + "\\Mutated_" + module.ModuleName);
                Console.WriteLine("****** Generated Mutant-Killing Constraint Number: "+Mutants.mutantNumber);

            }

        }
    }


    internal class HostEnvironment : MetadataReaderHost
    {
        PeReader peReader;
        public MyPlatformType myPlatformType;
        internal HostEnvironment()
        {
            this.peReader = new PeReader(this);
        }
        public override IUnit LoadUnitFrom(string location)
        {
            IUnit result = this.peReader.OpenModule(
              BinaryDocument.GetBinaryDocumentForFile(location, this));
            this.RegisterAsLatest(result);
            return result;
        }
        public MyPlatformType MyPlatformType
        {
            get
            {
                if (this.myPlatformType == null)
                {
                    this.myPlatformType = this.getMyPlatform();

                }
                return myPlatformType;
            }
        }
        protected MyPlatformType getMyPlatform()
        {
            return new MyPlatformType(this);
        }
    }

}
